/**
 * 
 */
package edu.gatech.cc.liam.core.rl.models.stochasticgame;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

import edu.gatech.cc.liam.core.DiscreteDistribution;
import edu.gatech.cc.liam.core.rl.models.Action;

/**
 * @author liam
 *
 */
public class DiscreteActionTransitionFunction<StateType> extends AbstractTransitionFunction<StateType> {
	
	private static final long serialVersionUID = -6545204281005405504L;
	
	//theProps maps States and then a jointAction to the distribution of next states
	final protected HashMap<StateType, HashMap<ArrayList<Action>, DiscreteDistribution<StateType>>> theProbs;
	final protected TreeMap<StateType, ArrayList<HashSet<Action>>> ActionSets;
	private int numberOfPlayers;
	
	public DiscreteActionTransitionFunction() {
		this.theProbs = new HashMap<StateType, HashMap<ArrayList<Action>,DiscreteDistribution<StateType>>>();
		ActionSets = new TreeMap<StateType, ArrayList<HashSet<Action>>>();
		numberOfPlayers = -1;
	}
	
	public void PutProb(StateType s1, ArrayList<Action> a, StateType s2,
			Double d) {
		PutProbHelper(s1, a, s2, d);
		
		ArrayList<HashSet<Action>> allPlayersActionsForGivenState = this.ActionSets.get(s1);
		if (allPlayersActionsForGivenState == null) {
			allPlayersActionsForGivenState = new ArrayList<HashSet<Action>>();
			ActionSets.put(s1, allPlayersActionsForGivenState);
			for(int i=0; i<a.size(); i++) {
				allPlayersActionsForGivenState.add(new HashSet<Action>());
			}
		}
		for(int i=0; i<a.size(); i++) {
			HashSet<Action> aPlayersActionSet = allPlayersActionsForGivenState.get(i);
			aPlayersActionSet.add(a.get(i));
		}
	}
	
	private void PutProbHelper(StateType s1, ArrayList<Action> a, StateType s2, Double d) {
		if(numberOfPlayers == -1) {
			numberOfPlayers = a.size();
		} else {
			assert(numberOfPlayers == a.size());
		}
		HashMap<ArrayList<Action>, DiscreteDistribution<StateType>> theStateProbs = theProbs.get(s1);
		if(theStateProbs == null) {
			theStateProbs = new HashMap<ArrayList<Action>, DiscreteDistribution<StateType>>();
			theProbs.put(s1, theStateProbs);
		}
		DiscreteDistribution<StateType> theDist;
		if(theStateProbs.containsKey(a)) {
			theDist = theStateProbs.get(a);
		} else {
			theDist = new DiscreteDistribution<StateType>();
			theStateProbs.put(a, theDist);
		}
		theDist.put(s2, d);
	}
	
	public DiscreteDistribution<StateType> TransitionDistribution(StateType oldState, ArrayList<Action> a) {
		HashMap<ArrayList<Action>, DiscreteDistribution<StateType>> theStateProbs = theProbs.get(oldState);
		assert(theStateProbs != null);
		return theStateProbs.get(a);
	}
	
	public double TransitionProb(StateType oldState, ArrayList<Action> a, StateType newState) {
		return TransitionDistribution(oldState,a).get(newState);
	}
	
	public boolean ValidateAllProbs() {
		for (StateType aState  : this.theProbs.keySet()) {
			HashMap<ArrayList<Action>, DiscreteDistribution<StateType>> actionProbsForAState = this.theProbs.get(aState);
			for (ArrayList<Action> aJointAction : actionProbsForAState.keySet()) {
				DiscreteDistribution<StateType> theDistribution = actionProbsForAState.get(aJointAction);
				if(theDistribution.isDistribution() == false) {
					System.out.println("**ERROR** The distribution " + theDistribution + " for state " + aState + " and joint action " + aJointAction + " is not a valid distribution!");
					System.out.println("Full Probabilities for state " + aState + ":\n" + actionProbsForAState);
					return false;
				}
			}
		}
		return true;
	}
	
	public void normalizeAllProbs() {
		for (Map.Entry<StateType, HashMap<ArrayList<Action>,DiscreteDistribution<StateType>>> e1 : this.theProbs.entrySet()) {
			for (Map.Entry<ArrayList<Action>, DiscreteDistribution<StateType>> e2 : e1.getValue().entrySet()) {
				e2.getValue().normalize();
			}
		}
	}
	
	@Override
	public String toString() {
		return "TransitionFunction [theProbs=" + theProbs + "]";
	}
	
	public Set<StateType> getStates() {
		return this.theProbs.keySet();
	}
	public Set<ArrayList<Action>> getAllJointsForState(StateType state) {
		return this.theProbs.get(state).keySet();
	}
	public int getNumberOfPlayers() {
		return this.numberOfPlayers;
	}
	public ArrayList<HashSet<Action>> getActionSetsForState(StateType theState) {
		return ActionSets.get(theState);
	}
	public Set<StateType> getSuccessorStateSet(StateType theState, ArrayList<Action> theJoint) {
		return this.theProbs.get(theState).get(theJoint).keySet();
	}
	
}
